package com.aiit.config;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;

import javax.websocket.OnClose;
import javax.websocket.OnMessage;
import javax.websocket.OnOpen;
import javax.websocket.Session;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Component
@ServerEndpoint("/ws")
public class WebSocketServer {
    private static final Logger LOGGER = LoggerFactory.getLogger(WebSocketServer.class);
    // 存储 username 到 Session 的映射
    private static final Map<String, Session> SESSION_USER_MAP = new ConcurrentHashMap<>();
    private static final int BUFFER_SIZE = 65536;

    @OnOpen
    public void onOpen(Session session) {
        try {
            session.setMaxBinaryMessageBufferSize(BUFFER_SIZE);
            session.setMaxTextMessageBufferSize(BUFFER_SIZE);
            // 暂时假设用户名通过其他方式获取，这里先用会话 ID 作为标识
            // 实际应用中可以在前端连接时发送用户名消息来设置
            SESSION_USER_MAP.put(session.getId(), session);
            LOGGER.info("新的 WebSocket 连接，会话 ID: {}", session.getId());
        } catch (Exception e) {
            LOGGER.error("WebSocket 连接打开时出错", e);
        }
    }

    // 处理文本消息
    @OnMessage
    public void onMessage(String message, Session session) {
        try {
            String sender = getUsernameBySession(session);
            LOGGER.info("接收到文本消息: {}，发送者: {}", message, sender);
            // 广播给其他用户（排除自己）
            broadcast(message, session);
        } catch (Exception e) {
            LOGGER.error("处理文本消息时出错", e);
        }
    }

    // 新增处理二进制消息（文件传输）
    @OnMessage
    public void onMessage(ByteBuffer byteBuffer, Session session) {
        try {
            String sender = getUsernameBySession(session);
            LOGGER.info("接收到二进制文件，发送者: {}", sender);
            // 广播二进制数据给其他用户（排除自己）
            broadcast(byteBuffer, session);
        } catch (Exception e) {
            LOGGER.error("处理二进制文件时出错", e);
            try {
                if (session.isOpen()) {
                    session.getBasicRemote().sendText("文件传输出错，请重试");
                }
            } catch (IOException ioException) {
                LOGGER.error("发送错误消息给客户端时出错", ioException);
            }
        }
    }

    @OnClose
    public void onClose(Session session) {
        try {
            String username = getUsernameBySession(session);
            if (username != null) {
                SESSION_USER_MAP.remove(username);
            }
            LOGGER.info("WebSocket 连接关闭，会话 ID: {}", session.getId());
        } catch (Exception e) {
            LOGGER.error("WebSocket 连接关闭时出错", e);
        }
    }

    // 广播文本消息给其他用户（排除发送者自己）
    public void broadcast(String message, Session excludeSession) {
        SESSION_USER_MAP.values().forEach(session -> {
            if (session.isOpen() && session != excludeSession) {
                try {
                    session.getBasicRemote().sendText(message);
                } catch (IOException e) {
                    LOGGER.error("文本消息广播错误: {}", e.getMessage());
                }
            }
        });
    }

    // 广播二进制消息给其他用户（排除发送者自己）
    private void broadcast(ByteBuffer byteBuffer, Session excludeSession) {
        byte[] data = new byte[byteBuffer.remaining()];
        byteBuffer.get(data);
        SESSION_USER_MAP.values().forEach(session -> {
            if (session.isOpen() && session != excludeSession) {
                try {
                    ByteBuffer bufferToSend = ByteBuffer.wrap(data);
                    session.getBasicRemote().sendBinary(bufferToSend);
                } catch (IOException e) {
                    LOGGER.error("二进制文件广播错误: {}", e.getMessage());
                }
            }
        });
    }

    // 新增：广播给所有商家（不排除任何人）
    public void broadcastToAll(String message) {
        SESSION_USER_MAP.values().forEach(session -> {
            if (session.isOpen()) {
                try {
                    session.getBasicRemote().sendText(message);
                } catch (IOException e) {
                    LOGGER.error("广播消息失败: {}", e.getMessage());
                }
            }
        });
    }

    private String getUsernameBySession(Session session) {
        for (Map.Entry<String, Session> entry : SESSION_USER_MAP.entrySet()) {
            if (entry.getValue().equals(session)) {
                return entry.getKey();
            }
        }
        return null;
    }
}